% This function simulate the discrete markov chain with exogenous exit.
% Input: 
%   TrMat (M x M):      
%       Transition probability of the state.
%   QW_0 (n x 1):
%       Histogram, or quadrature weight, of the initial distribution.
%   ExitProb (scaler between 0 and 1):
%       Probability of exogenous exit.
%   TrShockHist (N x T):
%       Uniformly distributed shocks for each sample point in each period.
%   ExitShockHist (N x T):
%       Uniformly distributed shocks for each sample point in each period.
% Output:
%   Sample (N x 2 x T):
%       In each period t, each observation n is presented by (state,age).
%       If the age is equal to 0, it indicates that this observation is an
%       entrant.
function Sample=MarkovChain_Simulation(TrMat,QW_0,ExitProb, ...
                                          TrShockHist,ExitShockHist)

%% Preliminary
M           =   size(TrMat,1);
% Check 1: whether TrMat is a square matrix
if size(TrMat,2)~=M
    error('Transition Matrix is not a square matrix');
end
% Check 2: dimension of QW_0
if size(QW_0,1)~=M
    error('Initial distribution dimension does not match');
end
% Check 3: dimension of ShockHist
[N,T]       =   size(TrShockHist);
if size(ExitShockHist,1)~=N || size(ExitShockHist,2)~=T
    error('Dimension of ExitShockHist is inconsistent with TrShockHist');
end

% Constants
AccTrMat    =   [zeros(1,M);cumsum(TrMat)];
AccQW_0     =   [0;cumsum(QW_0)];
if max(abs(AccQW_0(end)-1))>eps*10
    error('Initial QW is not summed up to 1');
end
if max(abs(AccTrMat(end,:)'-1))>eps*10
    error('Transition Matrix is not summed up to 1');
end
%% Initialization
Sample_0    =   [AssignLocation(AccQW_0,TrShockHist(:,1)),zeros(N,1)];
Sample      =   zeros(N,2,T);
Sample(:,:,1) ...
            =   Sample_0;
if T<=1
    return
end
%% Evolution 
for tt=2:T
    Ind_Exit    =   ( ExitShockHist(:,tt)<ExitProb );
    Ind_Exist   =   ~Ind_Exit;
    Age         =   zeros(N,1);
    ID          =   zeros(N,1);
    if any(Ind_Exit)
        Age(Ind_Exit)   =   0;
        TrShock_Exit    =   TrShockHist(Ind_Exit,tt);
        ID(Ind_Exit)    =   AssignLocation(AccQW_0,TrShock_Exit);
    end
    if any(Ind_Exist)
        
        LagID_Exist     =   Sample(Ind_Exist,1,tt-1);
        ID_Exist        =   zeros(length(LagID_Exist),1);
        TrShock_Exist   =   TrShockHist(Ind_Exist,tt);
        for nn=1:M
            TempInd     =   ( LagID_Exist==nn );
            if any(TempInd)
                ID_Exist(TempInd) ...
                            =   AssignLocation(AccTrMat(:,nn),...
                                               TrShock_Exist(TempInd));
            end
        end
        Age(Ind_Exist)  =   Sample(Ind_Exist,2,tt-1)+1;
        ID(Ind_Exist)   =   ID_Exist;
    end
    Sample(:,:,tt)= [ID,Age];
end

end
function ID=AssignLocation(AccQW,Shock)
    ID          =   celookup(AccQW,Shock);
    L           =   length(AccQW);
    ID(ID==0)   =   1;
    ID(ID==L)   =   L-1; 
end
